Skip to content

Support Autoscaling in Lakebase for Agents#353

Draft
jennsun wants to merge 4 commits intomainfrom
lakebase-autoscaling-support
Draft

Support Autoscaling in Lakebase for Agents#353
jennsun wants to merge 4 commits intomainfrom
lakebase-autoscaling-support

Conversation

@jennsun
Copy link
Contributor

@jennsun jennsun commented Mar 2, 2026

Add Lakebase autoscaling support (project/branch parameters) alongside the existing provisioned mode (instance_name) across all Lakebase classes and integrations

Using new autoscaling SDK: https://databricks-sdk-py.readthedocs.io/en/latest/workspace/postgres/postgres.html#databricks.sdk.service.postgres.PostgresAPI.list_endpoints

  • Autoscaling resolves hosts via postgres.list_endpoints (READ_WRITE endpoint) and mints tokens via postgres.generate_database_credential, while provisioned continues using database.get_database_instance and database.generate_database_credential
  • Strict parameter validation: exactly one mode must be specified — providing neither, partial (project without branch), or both provisioned + autoscaling params raises ValueError
  • Engine caching in AsyncDatabricksSession differentiates provisioned vs autoscaling cache keys
    (provisioned::{instance_name} vs autoscaling::{project}::{branch})
  • Comprehensive unit tests covering autoscaling host resolution, token minting, routing logic,
    validation edge cases, and engine caching across all four packages
  • update other classes using lakebasepool to support autoscaling

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant